import os,stat
import copy
import paramiko
from contextlib import contextmanager

RESULT_TEMPLATE = {
    "status": False,
    "stdout": "",
    "stderr": "",
    "error": "",
    "msg": ""
}


class SSHConnector:
    """SSH连接工具类"""

    def __init__(self, timeout=6):
        self.timeout = timeout
        
    def connect(self, hostname, username, password, port=22):
        """SSH连接"""
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(
            hostname=hostname,
            username=username,
            password=password,
            port=port,
            timeout=self.timeout,
            banner_timeout=20
        )
        return ssh

    def connect_close(self, ssh):
        """关闭SSH连接"""
        try:
            if ssh.get_transport() and ssh.get_transport().is_active():
                ssh.close()
        except:
            pass
    
    @contextmanager
    def ssh_session(self, server: dict):
        """上下文管理器"""
        ssh = self.connect(server["ip"], server["username"], server["password"])
        try:
            yield ssh 
        finally:
            self.connect_close(ssh)        

    def execute_command(self, server: dict, command: str, timeout=6):
        """执行命令"""
        result = copy.deepcopy(RESULT_TEMPLATE)
        try:
            with self.ssh_session(server) as ssh:  
                _, stdout, stderr = ssh.exec_command(command, timeout=timeout)
                stdout_lines = []
                stderr_lines = []

                for line in iter(stdout.readline, ""):
                    line = line.strip()
                    if not line:
                        continue
                    stdout_lines.append(line)
                err_output = stderr.read().decode().strip()
                if err_output:
                    stderr_lines.append(err_output)

                result["status"] = stdout.channel.recv_exit_status() == 0
                result["stdout"] = "\n".join(stdout_lines)
                result["stderr"] = "\n".join(stderr_lines)

        except (paramiko.SSHException, Exception) as e: 
            result["error"] = str(e)
        
        return result

    def upload_file(self, server: dict, local_path: str, remote_path: str):
        """上传文件"""
        result = copy.deepcopy(RESULT_TEMPLATE)
        try:
            with self.ssh_session(server) as ssh:
                sftp = ssh.open_sftp()
                remote_dir = os.path.dirname(remote_path)
                try:
                    sftp.stat(remote_dir)
                except IOError:
                    _, stdout, _ = ssh.exec_command(f"sudo mkdir -p {remote_dir}")
                    stdout.channel.recv_exit_status()

                sftp.put(local_path, remote_path)

                result["status"] = True
                result["msg"] = f"上传成功：{local_path} -> {remote_path}"
        except Exception as e: 
            result["error"] = str(e)
        finally:
            try:
                sftp.close()
            except:
                pass
        return result

    def upload_tar_and_extract(self, server: dict, local_tar_path: str, remote_dest_path: str):
        """上传压缩包并解压"""
        result = copy.deepcopy(RESULT_TEMPLATE)
        try:
            with self.ssh_session(server) as ssh:
                sftp = ssh.open_sftp()
                remote_tmp_tar = f"/tmp/{os.path.basename(local_tar_path)}"

                sftp.put(local_tar_path, remote_tmp_tar)
                cmd = f"sudo mkdir -p {remote_dest_path} && sudo tar -xzf {remote_tmp_tar} -C {remote_dest_path}"
                result= self.execute_command(server, cmd)
                
                sftp.remove(remote_tmp_tar)             
        except Exception as e: 
            result["error"] = str(e)
        finally:
            try:
                sftp.close()
            except:
                pass
        return result

    def download_with_tar(self, server: dict, remote_path: str, local_path: str):
        """打包下载远程目录"""
        result = copy.deepcopy(RESULT_TEMPLATE)
        try:
            with self.ssh_session(server) as ssh:
                sftp = ssh.open_sftp()
                tar_name = os.path.basename(remote_path.rstrip("/")) + ".tar.gz"
                remote_tmp_tar = f"/tmp/{tar_name}"
                local_tar_path = os.path.join(local_path, tar_name)

                # 执行远程打包命令
                cmd = f"sudo tar -czf {remote_tmp_tar} -C $(dirname {remote_path}) $(basename {remote_path})"
                cmd_result = self.execute_command(server, cmd)

                if not cmd_result.get("status"):
                    result.update(cmd_result)
                    result["msg"] = f"远程打包失败: {cmd_result.get('stderr') or cmd_result.get('stdout')}"
                    return result

                # 确认本地保存目录存在
                os.makedirs(local_path, exist_ok=True)

                # 下载文件
                try:
                    sftp.get(remote_tmp_tar, local_tar_path)
                    result["status"] = True
                    result["msg"] = local_tar_path
                except FileNotFoundError:
                    result["error"] = f"远程文件不存在: {remote_tmp_tar}"
                except Exception as e:
                    result["error"] = str(e)
                
                # 清理远程临时文件
                try:
                    sftp.remove(remote_tmp_tar)
                except:
                    pass

        except Exception as e:
            result["error"] = str(e)
        finally:
            try:
                sftp.close()
            except:
                pass
        return result

    def list_dir(self, server: dict, path: str):
        """列出远程目录内容，返回 list[dict]"""
        try:
            with self.ssh_session(server) as ssh:
                sftp = ssh.open_sftp()
                entries = []
                for f in sftp.listdir_attr(path):
                    entries.append({
                        "filename": f.filename,
                        "is_dir": stat.S_ISDIR(f.st_mode),
                        "size": f.st_size // 1024,
                        "mtime": f.st_mtime,
                        "mode": stat.filemode(f.st_mode)
                    })
                return {"status": True, "files": entries}
        except Exception as e:
            return {"status": False, "error": str(e)}

